"""
python p2_solution_simanneal_Annealer.py
"""

import json
import os
import random
import numpy as np
from simanneal import Annealer

# Paths for input and output files
warehouse_path = '../fujian/fujian3/origin_data/warehouse.json'
inventory_path = '../fujian/fujian3/data_from_p1/all_average_inventory.json'
sales_path = '../fujian/fujian3/data_from_p1/all_average_sales.json'
output_path = '../fujian/fujian3/simanneal_Annealer/'
os.makedirs(output_path, exist_ok=True)

# Load data
with open(warehouse_path, 'r') as f:
    warehouses = json.load(f)
with open(inventory_path, 'r') as f:
    inventory_data = json.load(f)
with open(sales_path, 'r') as f:
    sales_data = json.load(f)

# Process data into dictionaries for quick access
warehouse_dict = {w['warehouse_id']: w for w in warehouses}
inventory_dict = {i['category_id']: i['average_inventory'] for i in inventory_data}
sales_dict = {s['category_id']: s['average_sales'] for s in sales_data}

categories = list(inventory_dict.keys())
num_categories = len(categories)
warehouses_list = list(warehouse_dict.keys())
num_warehouses = len(warehouses_list)

# Define the Annealing class
class CategoryWarehouseAnnealer(Annealer):
    def __init__(self, initial_state):
        super().__init__(initial_state)

    def move(self):
        # Randomly reassign a category to a different warehouse
        category_idx = random.randint(0, num_categories - 1)
        new_warehouse = random.choice(warehouses_list)
        self.state[category_idx] = new_warehouse

    def energy(self):
        total_cost = 0
        warehouse_usage = {w: {'inventory': 0, 'sales': 0} for w in warehouses_list}

        # Calculate usage and cost
        for category_idx, warehouse_id in enumerate(self.state):
            category_id = categories[category_idx]
            inventory = inventory_dict[category_id]
            sales = sales_dict[category_id]
            warehouse = warehouse_dict[warehouse_id]

            # Update warehouse utilization
            warehouse_usage[warehouse_id]['inventory'] += inventory
            warehouse_usage[warehouse_id]['sales'] += sales

            # Add cost if this warehouse is used
            if inventory > 0:
                total_cost += warehouse['daily_cost']

        # Check constraints and penalize if constraints are violated
        penalty = 0
        for warehouse_id, usage in warehouse_usage.items():
            warehouse = warehouse_dict[warehouse_id]
            if usage['inventory'] > warehouse['max_inventory']:
                penalty += 1000 * (usage['inventory'] - warehouse['max_inventory'])
            if usage['sales'] > warehouse['max_sales']:
                penalty += 1000 * (usage['sales'] - warehouse['max_sales'])

        return total_cost + penalty

# Initialize random initial state (random warehouse assignments)
initial_state = [random.choice(warehouses_list) for _ in range(num_categories)]

# Run simulated annealing
annealer = CategoryWarehouseAnnealer(initial_state)
annealer.steps = 114514
best_state, best_cost = annealer.anneal()

# Prepare output based on best_state
warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in warehouses_list}
assignment = []

for category_idx, warehouse_id in enumerate(best_state):
    category_id = categories[category_idx]
    inventory = inventory_dict[category_id]
    sales = sales_dict[category_id]
    
    # Update utilization
    warehouse_utilization[warehouse_id]['inventory'] += inventory
    warehouse_utilization[warehouse_id]['sales'] += sales
    
    # Store assignment
    assignment.append({'category_id': category_id, 'warehouse_id': warehouse_id})

# Save assignment
with open(os.path.join(output_path, 'assignment.json'), 'w') as f:
    json.dump(assignment, f)

# Save total cost
with open(os.path.join(output_path, 'all_cost.txt'), 'w') as f:
    f.write(str(best_cost))

# Calculate and save utilization rates
utilization_rates_inventory = []
utilization_rates_sales = []

for warehouse_id, usage in warehouse_utilization.items():
    max_inventory = warehouse_dict[warehouse_id]['max_inventory']
    max_sales = warehouse_dict[warehouse_id]['max_sales']
    
    utilization_rates_inventory.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_inventory': usage['inventory'] / max_inventory
    })
    utilization_rates_sales.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_sales': usage['sales'] / max_sales
    })

# Save utilization rates to JSON files
with open(os.path.join(output_path, 'all_utilization_rate_inventory.json'), 'w') as f:
    json.dump(utilization_rates_inventory, f)
with open(os.path.join(output_path, 'all_utilization_rate_sales.json'), 'w') as f:
    json.dump(utilization_rates_sales, f)

print("Optimization Complete. Best total cost:", best_cost)
